package embedding

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/cloudwego/eino/components/embedding"
	"github.com/cloudwego/eino/schema"
	"github.com/patrickmn/go-cache"
	"github.com/sirupsen/logrus"
	"golang.org/x/time/rate"
)

// 使用示例和文档说明
/*
通用嵌入模型加载方法使用说明

本模块提供了基于模型名称的通用嵌入模型加载功能，支持多种主流嵌入模型。

主要功能：
1. 基于模型名称自动创建嵌入器
2. 支持OpenAI、智谱AI、Ollama、HuggingFace等多种模型
3. 自动模型类型检测
4. 配置合并和自定义
5. 模型注册表管理

使用方法：

1. 简单使用（推荐）
   embedder, err := NewEnhancedEmbedderFromModelName("text-embedding-ada-002")

2. 自定义配置
   customConfig := &EmbeddingConfig{
       APIKey: "your-api-key",
       BatchSize: 32,
       CacheEnabled: true,
   }
   embedder, err := CreateEmbedderFromModelName("text-embedding-3-small", customConfig)

3. 获取支持的模型列表
   registry := NewModelRegistry()
   models := registry.GetSupportedModels()

4. 自动检测模型类型
   modelType := AutoDetectModelType("text-embedding-ada-002") // 返回 "openai"

支持的模型：
- OpenAI: text-embedding-ada-002, text-embedding-3-small, text-embedding-3-large
- 智谱AI: embedding-2
- Ollama: nomic-embed-text, mxbai-embed-large
- HuggingFace: sentence-transformers/all-MiniLM-L6-v2, sentence-transformers/all-mpnet-base-v2
- 本地模型: local-bge-small, local-bge-base

配置优先级：
预定义模型配置 < 用户自定义配置
即用户提供的配置会覆盖预定义配置的相应字段。
*/

// Embedder 嵌入接口 - 保持不变以兼容现有代码
type Embedder interface {
	EmbedDocuments(ctx context.Context, docs []*schema.Document) error
	EmbedQuery(ctx context.Context, query string) ([]float64, error)
}

// EmbeddingConfig 嵌入配置结构体
type EmbeddingConfig struct {
	Type         string  `toml:"type" json:"type"`                   // 嵌入类型: local, openai, zhipuai, ollama, huggingface
	ModelName    string  `toml:"model_name" json:"model_name"`       // 模型名称
	ModelPath    string  `toml:"model_path" json:"model_path"`       // 本地模型路径
	EmbeddingDim int     `toml:"embedding_dim" json:"embedding_dim"` // 嵌入维度
	BatchSize    int     `toml:"batch_size" json:"batch_size"`       // 批处理大小
	MaxRetries   int     `toml:"max_retries" json:"max_retries"`     // 最大重试次数
	RetryDelay   int     `toml:"retry_delay" json:"retry_delay"`     // 重试延迟(秒)
	Timeout      int     `toml:"timeout" json:"timeout"`             // 超时时间(秒)
	CacheEnabled bool    `toml:"cache_enabled" json:"cache_enabled"` // 是否启用缓存
	CacheTTL     int     `toml:"cache_ttl" json:"cache_ttl"`         // 缓存过期时间(秒)
	APIKey       string  `toml:"api_key" json:"api_key"`             // API密钥
	BaseURL      string  `toml:"base_url" json:"base_url"`           // API基础URL
	Concurrency  int     `toml:"concurrency" json:"concurrency"`     // 并发数
	RateLimit    float64 `toml:"rate_limit" json:"rate_limit"`       // 速率限制(请求/秒)
}

// EmbeddingMetrics 嵌入指标
type EmbeddingMetrics struct {
	TotalRequests    int64         `json:"total_requests"`    // 总请求数
	SuccessRequests  int64         `json:"success_requests"`  // 成功请求数
	FailedRequests   int64         `json:"failed_requests"`   // 失败请求数
	CacheHits        int64         `json:"cache_hits"`        // 缓存命中数
	CacheMisses      int64         `json:"cache_misses"`      // 缓存未命中数
	AverageLatency   time.Duration `json:"average_latency"`   // 平均延迟
	TotalTokens      int64         `json:"total_tokens"`      // 总token数
	EmbeddingVectors int64         `json:"embedding_vectors"` // 生成的嵌入向量数
}

// EnhancedEmbedder 增强的嵌入器
type EnhancedEmbedder struct {
	config       EmbeddingConfig
	einoEmbedder embedding.Embedder
	cache        *cache.Cache
	metrics      *EmbeddingMetrics
	rateLimiter  *RateLimiter
	mu           sync.RWMutex
}

// RateLimiter 速率限制器，使用golang.org/x/time/rate实现
type RateLimiter struct {
	limiter *rate.Limiter
}

// NewRateLimiter 创建新的速率限制器
func NewRateLimiter(r float64) *RateLimiter {
	return &RateLimiter{
		limiter: rate.NewLimiter(rate.Limit(r), int(math.Ceil(r))), // 突发容量为速率的向上取整
	}
}

// ConfigManager 配置管理器，负责从文件加载和管理模型配置
type ConfigManager struct {
	configPath string
	config     EmbeddingConfigFile
	mu         sync.RWMutex
}

// EmbeddingConfigFile TOML配置文件结构体
type EmbeddingConfigFile struct {
	DefaultModel string                 `toml:"default_model"`
	Models       map[string]ModelConfig `toml:"models"`
}

// ModelRegistry 模型注册表，存储预定义的模型配置
type ModelRegistry struct {
	models        map[string]ModelConfig
	configManager *ConfigManager
}

// ModelConfig 模型配置信息
type ModelConfig struct {
	Type         string  `toml:"type" json:"type"`                   // 模型类型: local, openai, zhipuai, ollama, huggingface
	ModelName    string  `toml:"model_name" json:"model_name"`       // 模型名称
	EmbeddingDim int     `toml:"embedding_dim" json:"embedding_dim"` // 嵌入维度
	MaxTokens    int     `toml:"max_tokens" json:"max_tokens"`       // 最大token数
	APIKey       string  `toml:"api_key" json:"api_key"`             // API密钥
	BaseURL      string  `toml:"base_url" json:"base_url"`           // 基础URL
	Description  string  `toml:"description" json:"description"`     // 模型描述
	RateLimit    float64 `toml:"rate_limit" json:"rate_limit"`       // 速率限制
	ModelPath    string  `toml:"model_path" json:"model_path"`       // 本地模型路径
}

// ToEmbeddingConfig 将ModelConfig转换为EmbeddingConfig
func (mc *ModelConfig) ToEmbeddingConfig() EmbeddingConfig {
	return EmbeddingConfig{
		Type:         mc.Type,
		ModelName:    mc.ModelName,
		ModelPath:    mc.ModelPath,
		EmbeddingDim: mc.EmbeddingDim,
		APIKey:       mc.APIKey,
		BaseURL:      mc.BaseURL,
		RateLimit:    mc.RateLimit,
		// 设置默认值
		BatchSize:    16,
		MaxRetries:   3,
		RetryDelay:   1,
		Timeout:      30,
		CacheEnabled: true,
		CacheTTL:     3600,
		Concurrency:  4,
	}
}

// NewConfigManager 创建新的配置管理器
func NewConfigManager(configPath string) *ConfigManager {
	return &ConfigManager{
		configPath: configPath,
		config: EmbeddingConfigFile{
			Models: make(map[string]ModelConfig),
		},
	}
}

// LoadConfig 从文件加载配置
func (cm *ConfigManager) LoadConfig() error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	// 如果配置文件不存在，使用空配置
	if _, err := os.Stat(cm.configPath); os.IsNotExist(err) {
		logrus.Warnf("配置文件不存在: %s，将使用默认配置", cm.configPath)
		return nil
	}

	content, err := ioutil.ReadFile(cm.configPath)
	if err != nil {
		return fmt.Errorf("读取配置文件失败: %w", err)
	}

	if err := toml.Unmarshal(content, &cm.config); err != nil {
		return fmt.Errorf("解析配置文件失败: %w", err)
	}

	logrus.Infof("成功从 %s 加载配置，共 %d 个模型", cm.configPath, len(cm.config.Models))
	return nil
}

// GetModelConfig 从配置中获取模型配置
func (cm *ConfigManager) GetModelConfig(modelName string) (ModelConfig, bool) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	config, exists := cm.config.Models[modelName]
	return config, exists
}

// GetDefaultModel 获取默认模型名称
func (cm *ConfigManager) GetDefaultModel() string {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	return cm.config.DefaultModel
}

// NewModelRegistry 创建新的模型注册表
func NewModelRegistry(configPath ...string) *ModelRegistry {
	registry := &ModelRegistry{
		models: make(map[string]ModelConfig),
	}

	// 如果提供了配置文件路径，初始化配置管理器
	if len(configPath) > 0 && configPath[0] != "" {
		registry.configManager = NewConfigManager(configPath[0])
		if err := registry.configManager.LoadConfig(); err != nil {
			logrus.Warnf("加载配置文件失败: %v，将使用内置默认配置", err)
		}
	}

	return registry
}

// GetModelConfig 根据模型名称获取模型配置
func (r *ModelRegistry) GetModelConfig(modelName string) (ModelConfig, bool) {
	// 首先从配置文件中查找模型配置
	if r.configManager != nil {
		if config, exists := r.configManager.GetModelConfig(modelName); exists {
			logrus.Debugf("从配置文件中获取模型配置: %s", modelName)
			return config, true
		}
	}

	// 如果配置文件中没有，则从预定义模型中查找
	config, exists := r.models[modelName]
	return config, exists
}

// RegisterModel 注册新的模型配置
func (r *ModelRegistry) RegisterModel(modelName string, config ModelConfig) {
	r.models[modelName] = config
}

// GetSupportedModels 获取所有支持的模型列表
func (r *ModelRegistry) GetSupportedModels() []string {
	models := make([]string, 0, len(r.models))
	for name := range r.models {
		models = append(models, name)
	}
	return models
}

// AutoDetectModelType 自动检测模型类型
func AutoDetectModelType(modelName string) string {
	registry := NewModelRegistry()

	if modelConfig, exists := registry.GetModelConfig(modelName); exists {
		return modelConfig.Type
	}

	// 基于模型名称前缀进行简单检测
	if len(modelName) == 0 {
		return "local"
	}

	lowerName := strings.ToLower(modelName)

	if strings.Contains(lowerName, "openai") || strings.Contains(lowerName, "text-embedding") {
		return "openai"
	} else if strings.Contains(lowerName, "zhipu") || strings.Contains(lowerName, "embedding-2") {
		return "zhipuai"
	} else if strings.Contains(lowerName, "ollama") {
		return "ollama"
	} else if strings.Contains(lowerName, "sentence-transformers") || strings.Contains(lowerName, "huggingface") {
		return "huggingface"
	} else if strings.Contains(lowerName, "local") {
		return "local"
	}

	return "local" // 默认返回本地模型
}

// mergeModelConfigs 合并模型配置和用户自定义配置
func mergeModelConfigs(modelConfig ModelConfig, customConfig *EmbeddingConfig) EmbeddingConfig {
	// 使用ModelConfig的ToEmbeddingConfig方法获取基础配置
	result := modelConfig.ToEmbeddingConfig()

	// 如果用户提供了自定义配置，则覆盖相应字段
	if customConfig != nil {
		// 基本信息
		if customConfig.Type != "" {
			result.Type = customConfig.Type
		}
		if customConfig.ModelName != "" {
			result.ModelName = customConfig.ModelName
		}
		if customConfig.ModelPath != "" {
			result.ModelPath = customConfig.ModelPath
		}

		// 技术参数
		if customConfig.EmbeddingDim > 0 {
			result.EmbeddingDim = customConfig.EmbeddingDim
		}
		if customConfig.BatchSize > 0 {
			result.BatchSize = customConfig.BatchSize
		}
		if customConfig.Timeout > 0 {
			result.Timeout = customConfig.Timeout
		}
		if customConfig.Concurrency > 0 {
			result.Concurrency = customConfig.Concurrency
		}

		// 重试配置
		if customConfig.MaxRetries > 0 {
			result.MaxRetries = customConfig.MaxRetries
		}
		if customConfig.RetryDelay > 0 {
			result.RetryDelay = customConfig.RetryDelay
		}

		// 缓存配置
		result.CacheEnabled = customConfig.CacheEnabled
		if customConfig.CacheTTL > 0 {
			result.CacheTTL = customConfig.CacheTTL
		}

		// API配置
		if customConfig.APIKey != "" {
			result.APIKey = customConfig.APIKey
		}
		if customConfig.BaseURL != "" {
			result.BaseURL = customConfig.BaseURL
		}
		if customConfig.RateLimit > 0 {
			result.RateLimit = customConfig.RateLimit
		}
	}

	return result
}

// CreateEmbedderFromModelName 基于模型名称创建嵌入器
// configPath 可选的配置文件路径，如果不提供则使用默认配置
func CreateEmbedderFromModelName(modelName string, customConfig *EmbeddingConfig, configPath ...string) (*EnhancedEmbedder, error) {
	// 使用提供的配置文件路径或默认配置路径
	var registry *ModelRegistry
	if len(configPath) > 0 && configPath[0] != "" {
		registry = NewModelRegistry(configPath[0])
	} else {
		// 尝试使用默认配置文件路径
		defaultConfigPath := filepath.Join(".", "config", "embedding_models.toml")
		if _, err := os.Stat(defaultConfigPath); err == nil {
			registry = NewModelRegistry(defaultConfigPath)
		} else {
			registry = NewModelRegistry()
			logrus.Debug("未找到默认配置文件，使用内置配置")
		}
	}

	// 获取模型配置
	modelConfig, exists := registry.GetModelConfig(modelName)
	if !exists {
		return nil, fmt.Errorf("不支持的模型名称: %s", modelName)
	}

	// 合并配置
	finalConfig := mergeModelConfigs(modelConfig, customConfig)

	// 创建嵌入器
	return NewEnhancedEmbedder(finalConfig)
}

// 使用示例函数
func ExampleUsage() {
	fmt.Println("=== 通用嵌入模型加载方法使用示例 ===")

	// 示例1: 使用模型名称直接创建嵌入器（最简单的方式）
	fmt.Println("\n1. 使用模型名称直接创建嵌入器:")
	embedder1, err := NewEnhancedEmbedderFromModelName("text-embedding-ada-002")
	if err != nil {
		fmt.Printf("创建嵌入器失败: %v\n", err)
	} else {
		fmt.Println("成功创建OpenAI Ada 002嵌入器")
	}

	// 示例2: 使用自定义配置创建嵌入器
	fmt.Println("\n2. 使用自定义配置创建嵌入器:")
	customConfig := &EmbeddingConfig{
		APIKey:       "your-api-key",                  // 自定义API密钥
		BaseURL:      "https://custom-api.openai.com", // 自定义API地址
		BatchSize:    32,                              // 自定义批处理大小
		MaxRetries:   5,                               // 自定义重试次数
		CacheEnabled: true,                            // 启用缓存
		CacheTTL:     7200,                            // 缓存TTL: 2小时
	}

	embedder2, err := CreateEmbedderFromModelName("text-embedding-3-small", customConfig)
	if err != nil {
		fmt.Printf("创建嵌入器失败: %v\n", err)
	} else {
		fmt.Println("成功创建OpenAI Embedding 3 Small嵌入器，使用自定义配置")
	}

	// 示例3: 使用本地模型
	fmt.Println("\n3. 使用本地模型:")
	localConfig := &EmbeddingConfig{
		ModelPath: "./models/bge-small", // 本地模型路径
		BatchSize: 8,                    // 较小的批处理大小
	}

	embedder3, err := CreateEmbedderFromModelName("local-bge-small", localConfig)
	if err != nil {
		fmt.Printf("创建嵌入器失败: %v\n", err)
	} else {
		fmt.Println("成功创建本地BGE小型嵌入器")
	}

	// 示例4: 使用Ollama模型
	fmt.Println("\n4. 使用Ollama模型:")
	ollamaConfig := &EmbeddingConfig{
		BaseURL:   "http://localhost:11434", // Ollama服务地址
		BatchSize: 16,
	}

	embedder4, err := CreateEmbedderFromModelName("nomic-embed-text", ollamaConfig)
	if err != nil {
		fmt.Printf("创建嵌入器失败: %v\n", err)
	} else {
		fmt.Println("成功创建Ollama Nomic嵌入器")
	}

	// 示例5: 获取所有支持的模型列表
	fmt.Println("\n5. 获取支持的模型列表:")
	registry := NewModelRegistry()
	supportedModels := registry.GetSupportedModels()

	fmt.Println("支持的嵌入模型:")
	for _, model := range supportedModels {
		if config, exists := registry.GetModelConfig(model); exists {
			fmt.Printf("- %s: %s (维度: %d, 类型: %s)\n",
				model, config.Description, config.EmbeddingDim, config.Type)
		}
	}

	// 示例6: 自动检测模型类型
	fmt.Println("\n6. 自动检测模型类型:")
	testModels := []string{
		"text-embedding-ada-002",
		"embedding-2",
		"nomic-embed-text",
		"sentence-transformers/all-MiniLM-L6-v2",
		"custom-model",
	}

	for _, model := range testModels {
		modelType := AutoDetectModelType(model)
		fmt.Printf("模型 '%s' 检测为类型: %s\n", model, modelType)
	}

	// 避免未使用变量的编译错误
	_ = embedder1
	_ = embedder2
	_ = embedder3
	_ = embedder4
}

// Wait 等待直到有可用的令牌
func (rl *RateLimiter) Wait(ctx context.Context) error {
	// 使用golang.org/x/time/rate的Wait方法
	return rl.limiter.Wait(ctx)
}

// Allow 检查是否允许新的请求
func (rl *RateLimiter) Allow() bool {
	return rl.limiter.Allow()
}

// SetRate 设置新的速率限制
func (rl *RateLimiter) SetRate(r float64) {
	rl.limiter.SetLimit(rate.Limit(r))
}

// NewEnhancedEmbedder 创建新的增强嵌入器
func NewEnhancedEmbedder(config EmbeddingConfig) (*EnhancedEmbedder, error) {
	// 应用默认配置值
	config = applyDefaultConfig(config)

	// 自动检测模型类型（如果未指定）
	if config.Type == "" {
		config.Type = AutoDetectModelType(config.ModelName)
		logrus.Infof("自动检测到模型类型: %s", config.Type)
	}

	// 根据配置创建对应的嵌入器
	var einoEmbedder embedding.Embedder
	var err error

	switch config.Type {
	case "local":
		// 增强本地模型支持
		if config.ModelPath == "" {
			// 设置默认本地模型路径
			defaultModelDir := filepath.Join(".", "models")
			if config.ModelName != "" {
				config.ModelPath = filepath.Join(defaultModelDir, config.ModelName)
			} else {
				config.ModelPath = filepath.Join(defaultModelDir, "bge-base")
			}
			logrus.Infof("使用默认本地模型路径: %s", config.ModelPath)
		}

		// 验证本地模型路径
		if _, err := os.Stat(config.ModelPath); os.IsNotExist(err) {
			logrus.Warnf("本地模型路径不存在: %s，尝试下载或使用内置模型", config.ModelPath)
		}

		einoEmbedder, err = embedding.NewLocalEmbedder(&embedding.LocalEmbedderConfig{
			ModelPath:    config.ModelPath,
			EmbeddingDim: config.EmbeddingDim,
			BatchSize:    config.BatchSize,
		})

	case "openai":
		einoEmbedder, err = embedding.NewOpenAIEmbedder(&embedding.OpenAIEmbedderConfig{
			ModelName:    config.ModelName,
			EmbeddingDim: config.EmbeddingDim,
			APIKey:       config.APIKey,
			BaseURL:      config.BaseURL,
		})

	case "zhipuai":
		einoEmbedder, err = embedding.NewZhiPuAIEmbedder(&embedding.ZhiPuAIEmbedderConfig{
			ModelName:    config.ModelName,
			EmbeddingDim: config.EmbeddingDim,
			APIKey:       config.APIKey,
			BaseURL:      config.BaseURL,
		})

	case "ollama":
		// 为Ollama设置默认BaseURL
		if config.BaseURL == "" {
			config.BaseURL = "http://localhost:11434"
		}
		einoEmbedder, err = embedding.NewOllamaEmbedder(&embedding.OllamaEmbedderConfig{
			ModelName:    config.ModelName,
			EmbeddingDim: config.EmbeddingDim,
			BaseURL:      config.BaseURL,
		})

	case "huggingface":
		einoEmbedder, err = embedding.NewHuggingFaceEmbedder(&embedding.HuggingFaceEmbedderConfig{
			ModelName:    config.ModelName,
			EmbeddingDim: config.EmbeddingDim,
			APIKey:       config.APIKey,
		})

	default:
		return nil, fmt.Errorf("不支持的嵌入类型: %s", config.Type)
	}

	if err != nil {
		return nil, fmt.Errorf("创建嵌入器失败: %w", err)
	}

	// 初始化缓存
	var embeddingCache *cache.Cache
	if config.CacheEnabled {
		embeddingCache = cache.New(time.Duration(config.CacheTTL)*time.Second, 10*time.Minute)
		logrus.Infof("启用嵌入缓存，TTL: %d秒", config.CacheTTL)
	}

	// 初始化速率限制器
	var rateLimiter *RateLimiter
	if config.RateLimit > 0 {
		rateLimiter = NewRateLimiter(config.RateLimit)
		logrus.Infof("启用速率限制: %f请求/秒", config.RateLimit)
	}

	return &EnhancedEmbedder{
		config:       config,
		einoEmbedder: einoEmbedder,
		cache:        embeddingCache,
		metrics:      &EmbeddingMetrics{},
		rateLimiter:  rateLimiter,
		mu:           sync.RWMutex{},
	}, nil
}

// applyDefaultConfig 应用默认配置值
func applyDefaultConfig(config EmbeddingConfig) EmbeddingConfig {
	// 设置默认值（如果未指定）
	if config.BatchSize <= 0 {
		config.BatchSize = 16
	}
	if config.MaxRetries <= 0 {
		config.MaxRetries = 3
	}
	if config.RetryDelay <= 0 {
		config.RetryDelay = 1
	}
	if config.Timeout <= 0 {
		config.Timeout = 30
	}
	if config.CacheTTL <= 0 {
		config.CacheTTL = 3600
	}
	if config.Concurrency <= 0 {
		config.Concurrency = 4
	}

	// 为本地模型设置特殊默认值
	if config.Type == "local" && config.BatchSize > 8 {
		// 本地模型通常批处理大小较小
		config.BatchSize = 8
	}

	return config
}

// NewEnhancedEmbedderFromModelName 基于模型名称创建增强嵌入器（简化版本）
func NewEnhancedEmbedderFromModelName(modelName string) (*EnhancedEmbedder, error) {
	return CreateEmbedderFromModelName(modelName, nil)
}

// EmbedDocuments 为文档创建嵌入向量 - 实现Embedder接口
func (e *EnhancedEmbedder) EmbedDocuments(ctx context.Context, docs []*schema.Document) error {
	if len(docs) == 0 {
		logrus.Warn("文档列表为空")
		return nil
	}

	logrus.WithFields(logrus.Fields{
		"count":      len(docs),
		"batch_size": e.config.BatchSize,
		"model_type": e.config.Type,
		"model_name": e.config.ModelName,
	}).Info("开始为文档创建嵌入向量")

	startTime := time.Now()
	e.updateMetrics(func(m *EmbeddingMetrics) {
		m.TotalRequests++
		m.TotalTokens += int64(estimateTokens(docs))
	})

	// 按批次处理文档
	for i := 0; i < len(docs); i += e.config.BatchSize {
		end := i + e.config.BatchSize
		if end > len(docs) {
			end = len(docs)
		}

		batch := docs[i:end]
		if err := e.processBatch(ctx, batch); err != nil {
			e.updateMetrics(func(m *EmbeddingMetrics) {
				m.FailedRequests++
			})
			logrus.WithError(err).Error("处理文档批次失败")
			return fmt.Errorf("处理文档批次失败: %w", err)
		}

		logrus.WithFields(logrus.Fields{
			"processed": end,
			"total":     len(docs),
			"progress":  float64(end) / float64(len(docs)) * 100,
		}).Info("文档嵌入处理进度")
	}

	elapsed := time.Since(startTime)
	e.updateMetrics(func(m *EmbeddingMetrics) {
		m.SuccessRequests++
		m.EmbeddingVectors += int64(len(docs))
		m.AverageLatency = (m.AverageLatency*time.Duration(m.SuccessRequests-1) + elapsed) / time.Duration(m.SuccessRequests)
	})

	logrus.WithFields(logrus.Fields{
		"count":    len(docs),
		"duration": elapsed.Milliseconds(),
	}).Info("文档嵌入向量创建完成")

	return nil
}

// processBatch 处理文档批次
func (e *EnhancedEmbedder) processBatch(ctx context.Context, batch []*schema.Document) error {
	// 等待速率限制
	if e.rateLimiter != nil {
		for !e.rateLimiter.Allow() {
			time.Sleep(100 * time.Millisecond)
		}
	}

	// 提取文本内容
	batchTexts := make([]string, len(batch))
	for j, doc := range batch {
		batchTexts[j] = doc.Content
	}

	// 尝试获取嵌入向量（带重试机制）
	var embeddings [][]float64
	var err error

	for attempt := 0; attempt <= e.config.MaxRetries; attempt++ {
		if attempt > 0 {
			delay := time.Duration(e.config.RetryDelay) * time.Second * time.Duration(attempt)
			logrus.WithFields(logrus.Fields{
				"attempt": attempt,
				"delay":   delay,
			}).Warn("重试嵌入请求")
			time.Sleep(delay)
		}

		embeddings, err = e.einoEmbedder.EmbedDocuments(ctx, batchTexts)
		if err == nil {
			break
		}

		logrus.WithError(err).WithField("attempt", attempt).Error("嵌入请求失败")
	}

	if err != nil {
		return fmt.Errorf("嵌入请求最终失败: %w", err)
	}

	// 验证结果
	if len(embeddings) != len(batch) {
		return fmt.Errorf("嵌入向量数量与文档数量不匹配: %d != %d", len(embeddings), len(batch))
	}

	// 将嵌入向量分配给文档
	for j, doc := range batch {
		doc.Embedding = embeddings[j]

		// 缓存嵌入向量
		if e.cache != nil {
			cacheKey := generateCacheKey(doc.Content)
			e.cache.Set(cacheKey, embeddings[j], cache.DefaultExpiration)
		}
	}

	return nil
}

// EmbedQuery 为查询创建嵌入向量 - 实现Embedder接口
func (e *EnhancedEmbedder) EmbedQuery(ctx context.Context, query string) ([]float64, error) {
	if query == "" {
		return nil, fmt.Errorf("查询文本不能为空")
	}

	logrus.WithFields(logrus.Fields{
		"query_length": len(query),
		"model_type":   e.config.Type,
		"model_name":   e.config.ModelName,
	}).Debug("开始为查询创建嵌入向量")

	startTime := time.Now()
	e.updateMetrics(func(m *EmbeddingMetrics) {
		m.TotalRequests++
		m.TotalTokens += int64(estimateQueryTokens(query))
	})

	// 检查缓存
	if e.cache != nil {
		cacheKey := generateCacheKey(query)
		if cached, found := e.cache.Get(cacheKey); found {
			e.updateMetrics(func(m *EmbeddingMetrics) {
				m.SuccessRequests++
				m.CacheHits++
			})
			logrus.Debug("从缓存返回查询嵌入向量")
			return cached.([]float64), nil
		}
	}

	// 等待速率限制
	if e.rateLimiter != nil {
		for !e.rateLimiter.Allow() {
			time.Sleep(100 * time.Millisecond)
		}
	}

	// 尝试获取嵌入向量（带重试机制）
	var embedding []float64
	var err error

	for attempt := 0; attempt <= e.config.MaxRetries; attempt++ {
		if attempt > 0 {
			delay := time.Duration(e.config.RetryDelay) * time.Second * time.Duration(attempt)
			logrus.WithFields(logrus.Fields{
				"attempt": attempt,
				"delay":   delay,
			}).Warn("重试查询嵌入请求")
			time.Sleep(delay)
		}

		embedding, err = e.einoEmbedder.EmbedQuery(ctx, query)
		if err == nil {
			break
		}

		logrus.WithError(err).WithField("attempt", attempt).Error("查询嵌入请求失败")
	}

	if err != nil {
		e.updateMetrics(func(m *EmbeddingMetrics) {
			m.FailedRequests++
		})
		return nil, fmt.Errorf("查询嵌入请求最终失败: %w", err)
	}

	// 缓存结果
	if e.cache != nil {
		cacheKey := generateCacheKey(query)
		e.cache.Set(cacheKey, embedding, cache.DefaultExpiration)
	}

	elapsed := time.Since(startTime)
	e.updateMetrics(func(m *EmbeddingMetrics) {
		m.SuccessRequests++
		m.EmbeddingVectors++
		m.CacheMisses++
		m.AverageLatency = (m.AverageLatency*time.Duration(m.SuccessRequests-1) + elapsed) / time.Duration(m.SuccessRequests)
	})

	logrus.WithFields(logrus.Fields{
		"query_length":  len(query),
		"duration":      elapsed.Milliseconds(),
		"embedding_dim": len(embedding),
	}).Debug("查询嵌入向量创建完成")

	return embedding, nil
}

// GetMetrics 获取嵌入指标
func (e *EnhancedEmbedder) GetMetrics() EmbeddingMetrics {
	e.mu.RLock()
	defer e.mu.RUnlock()
	return *e.metrics
}

// ResetMetrics 重置指标
func (e *EnhancedEmbedder) ResetMetrics() {
	e.mu.Lock()
	defer e.mu.Unlock()
	e.metrics = &EmbeddingMetrics{}
}

// ClearCache 清除缓存
func (e *EnhancedEmbedder) ClearCache() {
	if e.cache != nil {
		e.cache.Flush()
		logrus.Info("嵌入缓存已清除")
	}
}

// GetCacheStats 获取缓存统计
func (e *EnhancedEmbedder) GetCacheStats() (items int, size int64) {
	if e.cache != nil {
		items = e.cache.ItemCount()
	}
	return items, size
}

// updateMetrics 更新指标（线程安全）
func (e *EnhancedEmbedder) updateMetrics(update func(*EmbeddingMetrics)) {
	e.mu.Lock()
	defer e.mu.Unlock()
	update(e.metrics)
}

// generateCacheKey 生成缓存键
func generateCacheKey(text string) string {
	// 简单的哈希实现，实际项目中可以使用更复杂的哈希算法
	return fmt.Sprintf("embedding_%s", text[:min(len(text), 50)])
}

// estimateTokens 估算token数量
func estimateTokens(docs []*schema.Document) int {
	total := 0
	for _, doc := range docs {
		// 粗略估算：中文字符算2个token，英文单词算1个token
		total += len(doc.Content) / 2
	}
	return total
}

// estimateQueryTokens 估算查询token数量
func estimateQueryTokens(query string) int {
	return len(query) / 2
}

// min 返回最小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// GenericEmbedder 通用嵌入器 - 兼容现有接口
type GenericEmbedder struct {
	enhancedEmbedder *EnhancedEmbedder
}

// NewGenericEmbedder 创建新的通用嵌入器（兼容现有接口）
func NewGenericEmbedder(embedType string, embeddingDim, batchSize int, modelPath string) (*GenericEmbedder, error) {
	config := EmbeddingConfig{
		Type:         embedType,
		EmbeddingDim: embeddingDim,
		BatchSize:    batchSize,
		ModelPath:    modelPath,
		MaxRetries:   3,
		RetryDelay:   1,
		Timeout:      30,
		CacheEnabled: true,
		CacheTTL:     3600,
		Concurrency:  4,
		RateLimit:    10,
	}

	enhancedEmbedder, err := NewEnhancedEmbedder(config)
	if err != nil {
		return nil, err
	}

	return &GenericEmbedder{
		enhancedEmbedder: enhancedEmbedder,
	}, nil
}

// EmbedDocuments 为文档创建嵌入向量
func (e *GenericEmbedder) EmbedDocuments(ctx context.Context, docs []*schema.Document) error {
	return e.enhancedEmbedder.EmbedDocuments(ctx, docs)
}

// EmbedQuery 为查询创建嵌入向量
func (e *GenericEmbedder) EmbedQuery(ctx context.Context, query string) ([]float64, error) {
	return e.enhancedEmbedder.EmbedQuery(ctx, query)
}

// ParallelEmbedder 并行嵌入器
type ParallelEmbedder struct {
	enhancedEmbedder *EnhancedEmbedder
	concurrency      int
}

// NewParallelEmbedder 创建新的并行嵌入器
func NewParallelEmbedder(embedder Embedder, concurrency int) *ParallelEmbedder {
	// 从通用嵌入器中提取增强嵌入器
	if genericEmbedder, ok := embedder.(*GenericEmbedder); ok {
		return &ParallelEmbedder{
			enhancedEmbedder: genericEmbedder.enhancedEmbedder,
			concurrency:      concurrency,
		}
	}

	// 如果传入的不是通用嵌入器，创建一个新的增强嵌入器
	config := EmbeddingConfig{
		Type:         "local",
		EmbeddingDim: 768,
		BatchSize:    16,
		MaxRetries:   3,
		RetryDelay:   1,
		Timeout:      30,
		CacheEnabled: true,
		CacheTTL:     3600,
		Concurrency:  concurrency,
		RateLimit:    10,
	}

	enhancedEmbedder, _ := NewEnhancedEmbedder(config)
	return &ParallelEmbedder{
		enhancedEmbedder: enhancedEmbedder,
		concurrency:      concurrency,
	}
}

// EmbedDocuments 并行为文档创建嵌入向量
func (pe *ParallelEmbedder) EmbedDocuments(ctx context.Context, docs []*schema.Document) error {
	if len(docs) == 0 {
		return nil
	}

	logrus.WithFields(logrus.Fields{
		"count":       len(docs),
		"concurrency": pe.concurrency,
	}).Info("开始并行处理文档嵌入")

	// 将文档分成多个批次
	batches := pe.splitDocumentsIntoBatches(docs, pe.concurrency)

	var wg sync.WaitGroup
	errChan := make(chan error, len(batches))

	// 启动多个协程处理批次
	for i, batch := range batches {
		if len(batch) == 0 {
			continue
		}

		wg.Add(1)
		go func(workerID int, batchDocs []*schema.Document) {
			defer wg.Done()

			logrus.WithFields(logrus.Fields{
				"worker_id":  workerID,
				"batch_size": len(batchDocs),
			}).Info("协程开始处理文档批次")

			if err := pe.enhancedEmbedder.EmbedDocuments(ctx, batchDocs); err != nil {
				logrus.WithError(err).WithField("worker_id", workerID).Error("协程处理文档嵌入失败")
				errChan <- err
				return
			}

			logrus.WithField("worker_id", workerID).Info("协程处理文档批次完成")
		}(i, batch)
	}

	// 等待所有协程完成
	go func() {
		wg.Wait()
		close(errChan)
	}()

	// 检查是否有错误
	for err := range errChan {
		if err != nil {
			return err
		}
	}

	logrus.Info("并行文档嵌入处理完成")
	return nil
}

// EmbedQuery 为查询创建嵌入向量
func (pe *ParallelEmbedder) EmbedQuery(ctx context.Context, query string) ([]float64, error) {
	return pe.enhancedEmbedder.EmbedQuery(ctx, query)
}

// splitDocumentsIntoBatches 将文档分成多个批次
func (pe *ParallelEmbedder) splitDocumentsIntoBatches(docs []*schema.Document, concurrency int) [][]*schema.Document {
	batches := make([][]*schema.Document, concurrency)
	for i := range batches {
		batches[i] = make([]*schema.Document, 0)
	}

	// 平均分配文档到各个批次
	for i, doc := range docs {
		batchIndex := i % concurrency
		batches[batchIndex] = append(batches[batchIndex], doc)
	}

	// 过滤空批次
	result := make([][]*schema.Document, 0, concurrency)
	for _, batch := range batches {
		if len(batch) > 0 {
			result = append(result, batch)
		}
	}

	return result
}
